// Copyright 2024 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package servertest

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"strings"
	"time"

	"golang.org/x/oauth2"

	"go.chromium.org/luci/auth/identity"
	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/data/stringset"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/grpc/grpcutil"

	"go.chromium.org/luci/server/auth"
)

// FakeRPCAuth implements auth.Method by accepting fake tokens as generated by
// FakeRPCTokenGenerator.
//
// It is used to authenticate incoming calls.
type FakeRPCAuth struct {
	// Seed is a random value expected to be found inside fake tokens.
	//
	// Each local server under test has its own seed. This is a precaution against
	// unrelated tests unexpectedly cross-talking to one another.
	Seed uint64

	// IDTokensAudience is ID token audiences to accept as valid.
	//
	// This normally includes "http[s]://127.0.0.1:<port>" addresses of the local
	// server, but may also include any other audiences passed via Options.
	//
	// If empty, ID tokens won't be accepted at all.
	IDTokensAudience stringset.Set
}

// Authenticate is part of auth.Method interface.
func (r *FakeRPCAuth) Authenticate(ctx context.Context, md auth.RequestMetadata) (*auth.User, auth.Session, error) {
	header := md.Header("Authorization")
	if header == "" {
		return nil, nil, nil // the auth method is not applicable
	}

	// Do not log the header in error messages below in case this is a real auth
	// token sent by mistake. We don't want to log real tokens in tests.

	parts := strings.Split(header, " ")
	if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
		return nil, nil, grpcutil.UnauthenticatedTag.Apply(errors.New("bad Authorization header, expecting \"Bearer ...\""))
	}

	blob, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, nil, grpcutil.UnauthenticatedTag.Apply(errors.Fmt("bad fake token, not base64-encoded: %w", err))
	}

	var tok FakeToken
	if err := json.Unmarshal(blob, &tok); err != nil {
		return nil, nil, grpcutil.UnauthenticatedTag.Apply(errors.Fmt("bad fake token, not JSON: %w", err))
	}

	if tok.Seed != r.Seed {
		return nil, nil,

			grpcutil.UnauthenticatedTag.Apply(errors.New("the fake token has unexpected seed, this usually means " +
				"independent tests are incorrectly talking to one another"))
	}

	switch tok.Kind {
	case "oauth":
		// Any scopes set in a fake OAuth token is fine for now.
	case "id":
		if len(r.IDTokensAudience) == 0 {
			return nil, nil,

				grpcutil.UnauthenticatedTag.Apply(errors.Fmt("unexpectedly got a fake ID token with audience %q, "+
					"the server is not configured to accept ID tokens at all", tok.Aud))
		}
		if !r.IDTokensAudience.Has(tok.Aud) {
			return nil, nil,
				grpcutil.UnauthenticatedTag.Apply(errors.Fmt("unexpected ID token audience %q, expecting one of %v", tok.Aud, r.IDTokensAudience.ToSortedSlice()))
		}
	default:
		return nil, nil, grpcutil.UnauthenticatedTag.Apply(errors.Fmt("unexpected fake token kind %q", tok.Kind))
	}

	exp := time.Unix(tok.Expiry, 0)
	if dt := clock.Now(ctx).Sub(exp); dt > 0 {
		return nil, nil, grpcutil.UnauthenticatedTag.Apply(errors.Fmt("fake token expired %s ago", dt))
	}

	ident, err := identity.MakeIdentity("user:" + tok.Email)
	if err != nil {
		return nil, nil, grpcutil.UnauthenticatedTag.Apply(errors.Fmt("malformed email %q in the fake token", tok.Email))
	}

	return &auth.User{
		Identity: ident,
		Email:    tok.Email,
		Extra:    &tok,
	}, nil, nil
}

// GetFakeToken returns an info about the fake token used to authenticate this
// call or nil if this call wasn't authenticated by a fake token.
func GetFakeToken(ctx context.Context) *FakeToken {
	tok, _ := auth.CurrentUser(ctx).Extra.(*FakeToken)
	return tok
}

// FakeRPCTokenGenerator knows how to produce tokens understood by FakeRPCAuth.
//
// Implements localauth.TokenGenerator.
type FakeRPCTokenGenerator struct {
	// Seed is a random value to put into fake tokens.
	//
	// Each local server under test has its own seed. This is a precaution against
	// unrelated tests unexpectedly cross-talking to one another.
	Seed uint64

	// Email to put into fake tokens.
	//
	// It will show up as an authenticated user email on the server side.
	Email string
}

// GenerateOAuthToken is part of localauth.TokenGenerator interface.
func (g *FakeRPCTokenGenerator) GenerateOAuthToken(ctx context.Context, scopes []string, _ time.Duration) (*oauth2.Token, error) {
	return genFakeToken(ctx, &FakeToken{
		Kind:   "oauth",
		Email:  g.Email,
		Seed:   g.Seed,
		Scopes: scopes,
	})
}

// GenerateIDToken is part of localauth.TokenGenerator interface.
func (g *FakeRPCTokenGenerator) GenerateIDToken(ctx context.Context, audience string, _ time.Duration) (*oauth2.Token, error) {
	return genFakeToken(ctx, &FakeToken{
		Kind:  "id",
		Email: g.Email,
		Seed:  g.Seed,
		Aud:   audience,
	})
}

// GetEmail is part of localauth.TokenGenerator interface.
func (g *FakeRPCTokenGenerator) GetEmail() (string, error) {
	return g.Email, nil
}

// genFakeToken generates a base64-encoded fake token containing given info.
func genFakeToken(ctx context.Context, tok *FakeToken) (*oauth2.Token, error) {
	expiry := clock.Now(ctx).Add(15 * time.Minute).Round(time.Second)
	tok.Expiry = expiry.Unix()
	blob, err := json.Marshal(tok)
	if err != nil {
		return nil, err // should be impossible, we marshal simple types
	}
	return &oauth2.Token{
		AccessToken: base64.RawURLEncoding.EncodeToString(blob),
		Expiry:      expiry,
	}, nil
}

// FakeToken is a payload of a fake RPC auth token.
type FakeToken struct {
	Kind   string   `json:"kind"`
	Email  string   `json:"email"`
	Expiry int64    `json:"expiry"`
	Seed   uint64   `json:"seed"`
	Aud    string   `json:"aud,omitempty"`
	Scopes []string `json:"scopes,omitempty"`
}
